import os
import json
import numpy as np
import logging
from typing import List, Dict, Tuple
import pickle
import re

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class SimpleRAGSystem:
    """简化版RAG系统 - 不依赖sentence-transformers"""
    
    def __init__(self, knowledge_base_path='knowledge_base'):
        """
        初始化简化版RAG系统
        
        Args:
            knowledge_base_path: 知识库存储路径
        """
        self.knowledge_base_path = knowledge_base_path
        self.knowledge_base = []
        self.knowledge_file = os.path.join(knowledge_base_path, 'knowledge_base.json')
        
        # 创建知识库目录
        os.makedirs(knowledge_base_path, exist_ok=True)
        
        # 加载或创建知识库
        self._load_knowledge_base()
    
    def _load_knowledge_base(self):
        """加载知识库"""
        try:
            if os.path.exists(self.knowledge_file):
                logger.info("加载现有知识库...")
                with open(self.knowledge_file, 'r', encoding='utf-8') as f:
                    self.knowledge_base = json.load(f)
                logger.info(f"知识库加载完成，包含 {len(self.knowledge_base)} 条知识")
            else:
                logger.info("创建新的知识库...")
                self._create_default_knowledge_base()
        except Exception as e:
            logger.error(f"加载知识库失败: {e}")
            self._create_default_knowledge_base()
    
    def _create_default_knowledge_base(self):
        """创建默认的乳腺癌知识库"""
        default_knowledge = [
            {
                "id": 1,
                "title": "乳腺癌概述",
                "content": "乳腺癌是女性最常见的恶性肿瘤之一，起源于乳腺上皮细胞。早期发现和治疗可以显著提高治愈率。全球每年约有200万新发病例，是女性癌症死亡的主要原因之一。",
                "category": "基础概念"
            },
            {
                "id": 2,
                "title": "乳腺癌症状",
                "content": "乳腺癌常见症状包括：乳房肿块（通常无痛、质地硬、边界不清）、乳头溢液（特别是血性溢液）、乳房皮肤改变（如橘皮样改变、凹陷）、乳头内陷、腋窝淋巴结肿大等。早期乳腺癌可能无明显症状。",
                "category": "症状"
            },
            {
                "id": 3,
                "title": "乳腺癌诊断方法",
                "content": "乳腺癌诊断方法包括：乳腺X线摄影（钼靶）、乳腺超声、乳腺MRI、活检等。细胞学检查是确诊的重要手段。影像学检查可以发现可疑病变，但最终确诊需要病理检查。",
                "category": "诊断"
            },
            {
                "id": 4,
                "title": "乳腺癌分期",
                "content": "乳腺癌分期采用TNM分期系统：T表示肿瘤大小（T1-T4），N表示淋巴结转移（N0-N3），M表示远处转移（M0-M1）。分为0期（原位癌）到IV期（远处转移）。分期越早，预后越好。",
                "category": "分期"
            },
            {
                "id": 5,
                "title": "乳腺癌治疗方法",
                "content": "乳腺癌治疗方法包括：手术治疗（保乳手术、全乳切除术）、化疗、放疗、内分泌治疗、靶向治疗等。治疗方案需根据患者具体情况制定，包括肿瘤分期、分子分型、患者年龄和身体状况。",
                "category": "治疗"
            },
            {
                "id": 6,
                "title": "乳腺癌预防",
                "content": "乳腺癌预防措施包括：定期乳腺检查（自检、临床检查、影像学检查）、保持健康生活方式（规律运动、健康饮食、控制体重）、避免长期使用雌激素、遗传咨询等。高危人群应加强监测。",
                "category": "预防"
            },
            {
                "id": 7,
                "title": "乳腺癌预后",
                "content": "乳腺癌预后与分期、分子分型、治疗反应等因素相关。早期发现和治疗的患者5年生存率可达90%以上。I期患者5年生存率超过95%，IV期患者约为25%。",
                "category": "预后"
            },
            {
                "id": 8,
                "title": "乳腺癌分子分型",
                "content": "乳腺癌分子分型包括：Luminal A型（激素受体阳性，HER2阴性，Ki67低）、Luminal B型（激素受体阳性，HER2阳性或Ki67高）、HER2阳性型、三阴性型（激素受体和HER2均阴性）。不同分型需要不同的治疗方案。",
                "category": "分子分型"
            },
            {
                "id": 9,
                "title": "乳腺癌筛查",
                "content": "乳腺癌筛查包括：乳腺自检（每月一次）、临床乳腺检查（每年一次）、乳腺X线摄影（40岁以上女性每年一次）。高危人群可能需要更频繁的筛查，包括MRI检查。",
                "category": "筛查"
            },
            {
                "id": 10,
                "title": "乳腺癌遗传因素",
                "content": "乳腺癌遗传因素包括BRCA1、BRCA2基因突变等。有家族史的女性应进行遗传咨询和基因检测。BRCA1/2突变携带者一生中患乳腺癌的风险高达60-80%。",
                "category": "遗传"
            }
        ]
        
        self.knowledge_base = default_knowledge
        self._save_knowledge_base()
    
    def _save_knowledge_base(self):
        """保存知识库到文件"""
        try:
            with open(self.knowledge_file, 'w', encoding='utf-8') as f:
                json.dump(self.knowledge_base, f, ensure_ascii=False, indent=2)
            logger.info("知识库保存成功")
        except Exception as e:
            logger.error(f"保存知识库失败: {e}")
    
    def _simple_similarity(self, query: str, text: str) -> float:
        """简单的文本相似度计算"""
        query_words = set(re.findall(r'\w+', query.lower()))
        text_words = set(re.findall(r'\w+', text.lower()))
        
        if not query_words:
            return 0.0
        
        intersection = query_words.intersection(text_words)
        return len(intersection) / len(query_words)
    
    def search(self, query: str, top_k: int = 3) -> List[Dict]:
        """
        搜索相关知识（简化版）
        
        Args:
            query: 查询文本
            top_k: 返回最相关的前k条知识
            
        Returns:
            相关知识的列表
        """
        try:
            if not self.knowledge_base:
                logger.warning("知识库为空，无法搜索")
                return []
            
            # 计算相似度
            results = []
            for item in self.knowledge_base:
                similarity = self._simple_similarity(query, item['content'])
                if similarity > 0.1:  # 降低阈值
                    knowledge = item.copy()
                    knowledge['similarity'] = similarity
                    results.append(knowledge)
            
            # 按相似度排序
            results.sort(key=lambda x: x['similarity'], reverse=True)
            
            logger.info(f"搜索完成，找到 {len(results)} 条相关知识")
            return results[:top_k]
            
        except Exception as e:
            logger.error(f"搜索失败: {e}")
            return []
    
    def get_answer(self, question: str, context_limit: int = 2) -> Dict:
        """
        基于RAG生成答案（简化版）
        
        Args:
            question: 用户问题
            context_limit: 使用的上下文知识条数
            
        Returns:
            包含答案和相关知识的字典
        """
        try:
            # 搜索相关知识
            relevant_knowledge = self.search(question, top_k=context_limit)
            
            if not relevant_knowledge:
                return {
                    "answer": "抱歉，我没有找到相关的医学知识来回答您的问题。建议您咨询专业医生。",
                    "sources": [],
                    "confidence": 0.0
                }
            
            # 构建答案
            answer_parts = []
            sources = []
            
            for knowledge in relevant_knowledge:
                answer_parts.append(f"根据医学知识：{knowledge['content']}")
                sources.append({
                    "title": knowledge['title'],
                    "category": knowledge['category'],
                    "similarity": knowledge['similarity']
                })
            
            answer = " ".join(answer_parts)
            
            # 计算置信度（基于最高相似度）
            confidence = max([k['similarity'] for k in relevant_knowledge])
            
            return {
                "answer": answer,
                "sources": sources,
                "confidence": confidence
            }
            
        except Exception as e:
            logger.error(f"生成答案失败: {e}")
            return {
                "answer": "抱歉，系统出现错误，无法生成答案。",
                "sources": [],
                "confidence": 0.0
            }
    
    def add_knowledge(self, title: str, content: str, category: str = "其他"):
        """添加新知识到知识库"""
        try:
            new_id = max([item['id'] for item in self.knowledge_base], default=0) + 1
            new_knowledge = {
                "id": new_id,
                "title": title,
                "content": content,
                "category": category
            }
            
            self.knowledge_base.append(new_knowledge)
            self._save_knowledge_base()
            
            logger.info(f"新知识添加成功: {title}")
            return True
        except Exception as e:
            logger.error(f"添加知识失败: {e}")
            return False
    
    def get_knowledge_by_category(self, category: str) -> List[Dict]:
        """根据类别获取知识"""
        return [item for item in self.knowledge_base if item['category'] == category]
    
    def get_all_categories(self) -> List[str]:
        """获取所有知识类别"""
        categories = set(item['category'] for item in self.knowledge_base)
        return list(categories)
    
    def get_knowledge_stats(self) -> Dict:
        """获取知识库统计信息"""
        try:
            total_knowledge = len(self.knowledge_base)
            categories = self.get_all_categories()
            category_counts = {}
            
            for item in self.knowledge_base:
                category = item['category']
                category_counts[category] = category_counts.get(category, 0) + 1
            
            return {
                'total_knowledge': total_knowledge,
                'total_categories': len(categories),
                'categories': categories,
                'category_counts': category_counts
            }
        except Exception as e:
            logger.error(f"获取知识库统计信息失败: {e}")
            return {
                'total_knowledge': 0,
                'total_categories': 0,
                'categories': [],
                'category_counts': {}
            }
